'''
Created on Sep 28, 2009

@author: mkiyer
'''
from veggie.genome.chrom import get_chrom_names
import logging
import sys

__refflat_path = "/lab/mkiyer/refdb/ucsc_tables/refFlat.txt"

class Gene(object):
    '''
    Stores relevant information about genes.  Eventually this should be
    merged into veggie as a gene annotation database rather than read
    from a flat file each time the program is run
    
    Currently this is just a 'container' object with no purpose other
    than to group gene properties together into a single object
    '''
    def __init__(self):
        '''
        for now leave this empty and let other methods assign gene 
        properties arbitrarily
        '''
        pass
    
    def is_protein_coding(self):
        return self.cds_start != self.cds_end

def parse_refseq_genes():
    '''
    Generates Gene objects from the refFlat.txt gene annotation file
    
    This should be merged into veggie eventually, so that all programs that
    need to use the refseq genes can do so with the same code.
    
    refFlat.txt has the following schema:
    -------------------------------------
    hugo name varchar(255) NOT NULL default '',
    name varchar(255) NOT NULL default '',
    chrom varchar(255) NOT NULL default '',
    strand char(1) NOT NULL default '',
    txStart int(10) unsigned NOT NULL default '0',
    txEnd int(10) unsigned NOT NULL default '0',
    cdsStart int(10) unsigned NOT NULL default '0',
    cdsEnd int(10) unsigned NOT NULL default '0',
    exonCount int(10) unsigned NOT NULL default '0',
    exonStarts longblob NOT NULL,
    exonEnds longblob NOT NULL,
    proteinID varchar(40) NOT NULL default '',
    alignID varchar(255) NOT NULL default '',
    '''
    for line in open(__refflat_path):
        fields = line.strip().split('\t')
        g = Gene()
        g.symbol = fields[0]
        g.acc = fields[1]
        g.chrom = fields[2]
        if g.chrom not in get_chrom_names():
            continue
        g.strand = fields[3]
        if g.strand is not '+' and g.strand is not '-':
            raise ValueError("gene strand must be either + or -")
        
        g.tx_start = int(fields[4])
        g.tx_end = int(fields[5])
        g.cds_start = int(fields[6])
        g.cds_end = int(fields[7])
        g.exon_count = int(fields[8])
        g.exon_starts = map(int, fields[9].split(',')[:-1])
        g.exon_ends = map(int, fields[10].split(',')[:-1])
        g.exons = zip(g.exon_starts, g.exon_ends)
        g.introns = zip(g.exon_ends, g.exon_starts[1:])
        
        # build a list of all the 3'UTR exons
        utr5_exons = []
        utr3_exons = []
        for e in g.exons:
            if e[0] < g.cds_start:
                utr5_exons.append(e)                
            if e[1] > g.cds_end:
                utr3_exons.append(e)
        if g.strand == '+':
            g.utr5_exons = utr5_exons
            g.utr3_exons = utr3_exons
        else:
            g.utr5_exons = utr3_exons
            g.utr3_exons = utr5_exons

        # find first/last coding exons
        first_cds_exon = (g.tx_end, g.tx_end)
        last_cds_exon = (g.tx_end, g.tx_end)
        if g.tx_start != g.cds_start:
            for e in g.exons[::-1]:
                first_cds_exon = (g.cds_start, e[1])
                if e[0] <= g.cds_start:
                    break
        if g.tx_end != g.cds_end:
            for e in g.exons:
                last_cds_exon = (e[0], g.cds_end)
                if e[1] >= g.cds_end:
                    break
        if g.strand == '-':
            g.utr5_cds = last_cds_exon
            g.utr3_cds = first_cds_exon
        else:
            g.utr5_cds = first_cds_exon
            g.utr3_cds = last_cds_exon

        # find the "noncoding" parts of the UTR
        g.utr3 = (g.tx_end, g.tx_end)        
        if g.tx_start != g.cds_start:
            if g.strand == '+':
                g.utr5 = (g.tx_start, g.cds_start)
            else:
                g.utr3 = (g.tx_start, g.cds_start)
        if g.tx_end != g.cds_end:
            if g.strand == '+':
                g.utr3 = (g.cds_end, g.tx_end)
            else:
                g.utr5 = (g.cds_end, g.tx_end)

        yield g

        # the entire 3'UTR interval (not just the exons)
        # it is possible for a 3'UTR to have multiple exons
        # this interval includes all the exons and introns
        # leading to the tx_end
        #g.utr3_interval = min([e[0] for e in g.utr3_exons]), max([e[1] for e in g.utr3_exons])

def parse_bed_file(fhd):
    '''
    parse a gene bed file
    '''
    for linenum, line in enumerate(fhd):
        if line is None:
            continue
        line = line.strip()
        if line.startswith('#'):
            logging.debug("skipping comment line %d: %s" % (linenum, line))
            continue
        if line.startswith('track'):
            logging.debug("skipping track header line %d: %s"  % (linenum, line))
            continue
        thisfields = line.split()
        # first three fields are required
        g = Gene()
        g.chrom = thisfields[0]
        g.tx_start = int(thisfields[1])
        g.tx_end = int(thisfields[2])
        g.symbol = thisfields[3]
        g.strand = thisfields[5]
        g.cds_start = int(thisfields[6])
        g.cds_end = int(thisfields[7])
        g.exon_count = int(thisfields[9])
        block_sizes = map(int, thisfields[10].split(',')[:-1])
        block_starts = map(int, thisfields[11].split(',')[:-1])        
        g.exon_starts = [(g.tx_start + start) for start in block_starts]        
        g.exon_ends = [(start + size) for start, size in zip(g.exon_starts, block_sizes)]
        g.introns = zip(g.exon_ends, g.exon_starts[1:])        
        yield g    


def parse_gene_pred(fhd):
    '''
    name varchar(255) NOT NULL default '',
    chrom varchar(255) NOT NULL default '',
    strand char(1) NOT NULL default '',
    txStart int(10) unsigned NOT NULL default '0',
    txEnd int(10) unsigned NOT NULL default '0',
    cdsStart int(10) unsigned NOT NULL default '0',
    cdsEnd int(10) unsigned NOT NULL default '0',
    exonCount int(10) unsigned NOT NULL default '0',
    exonStarts longblob NOT NULL,
    exonEnds longblob NOT NULL,
    proteinID varchar(40) NOT NULL default '',
    alignID varchar(255) NOT NULL default '',
    '''
    for line in fhd:
        if line.startswith('#'):
            continue
        fields = line.strip().split('\t')
        name = fields[0]
        chrom = fields[1]
        strand = fields[2]
        txstart = int(fields[3])
        txend = int(fields[4])
        yield (name, txstart, txend)
